import numpy as np
import math
import pygad

# -------------------------- 1. 基础参数与复用函数（严格遵循A题.pdf） --------------------------
target_true = np.array([0, 200, 0])  # 真目标底面圆心（文档1-6）
r, h = 7, 10  # 真目标半径7m、高10m（文档1-6）
target_fake = np.array([0, 0, 0])  # 假目标（文档1-6）
missile_M1 = np.array([20000, 0, 2000])  # M1初始位置（文档1-6）
missile_speed = 300  # 导弹速度300m/s（文档1-6）
drone_FY1 = np.array([17800, 0, 1800])  # FY1初始位置（文档1-6）
g = 9.8  # 重力加速度（文档1-5）
smoke_radius = 10  # 烟幕有效半径10m（文档1-5）
smoke_sink_speed = 3  # 烟幕下沉速度3m/s（文档1-5）
smoke_valid_time = 20  # 烟幕起爆后20s有效（文档1-5）
drone_speed_range = [70, 140]  # 无人机速度约束（文档1-7）
NUM_SMOKE = 3  # 第三题核心：3枚烟幕弹（文档1-11）

# 计算导弹飞到假目标的时间（关键约束，文档1-6）
missile_distance_to_fake = np.linalg.norm(target_fake - missile_M1)
missile_time_to_fake = missile_distance_to_fake / missile_speed  # ≈67.0s


# 1.1 真目标外接长方体顶点生成（文档1-5“有效遮蔽需覆盖真目标”）
def get_cylinder_bounding_box_vertices(target_center, radius, height):
    x_min = target_center[0] - radius
    x_max = target_center[0] + radius
    y_min = target_center[1] - radius
    y_max = target_center[1] + radius
    z_min = target_center[2]
    z_max = target_center[2] + height
    vertices = [
        np.array([x_min, y_min, z_min]), np.array([x_max, y_min, z_min]),
        np.array([x_max, y_max, z_min]), np.array([x_min, y_max, z_min]),
        np.array([x_min, y_min, z_max]), np.array([x_max, y_min, z_max]),
        np.array([x_max, y_max, z_max]), np.array([x_min, y_max, z_max])
    ]
    return vertices


# 预计算目标顶点（减少冗余计算）
target_vertices = get_cylinder_bounding_box_vertices(target_true, r, h)


# 1.2 导弹实时位置计算（文档1-6“直指假目标飞行”）
def missile_real_time_position(missile_init_pos, t):
    if t >= missile_time_to_fake:  # 导弹命中假目标后停止（文档1-6）
        return target_fake
    if t < 0:
        return missile_init_pos
    dir_vec = target_fake - missile_init_pos
    dir_norm = np.linalg.norm(dir_vec)
    if dir_norm <= 1e-6:
        return missile_init_pos
    dir_unit_vec = dir_vec / dir_norm
    return missile_init_pos + missile_speed * t * dir_unit_vec


# 1.3 烟幕起爆点计算（文档1-5“重力作用+时序起爆”）
def get_smoke_explosion_center(drone_init_pos, drone_speed, theta, t_drop, t_detonate):
    # 投放点坐标（无人机t_drop时刻位置）
    drop_x = drone_init_pos[0] + drone_speed * math.cos(theta) * t_drop
    drop_y = drone_init_pos[1] + drone_speed * math.sin(theta) * t_drop
    drop_z = drone_init_pos[2]
    # 烟幕弹平抛运动（水平同无人机速度，竖直自由下落）
    explosion_x = drop_x + drone_speed * math.cos(theta) * t_detonate
    explosion_y = drop_y + drone_speed * math.sin(theta) * t_detonate
    fall_distance = 0.5 * g * (t_detonate ** 2)
    explosion_z = max(drop_z - fall_distance, 10.0)  # 避免落地（文档1-5隐含约束）
    return np.array([explosion_x, explosion_y, explosion_z])


# 1.4 烟幕实时位置计算（文档1-5“匀速下沉”）
def smoke_real_time_center(drone_init_pos, drone_speed, theta, t_drop, t_detonate, t):
    detonate_t = t_drop + t_detonate
    if t < detonate_t - 1e-6:  # 起爆前烟幕未形成
        return np.array([0, 0, 0])
    explosion_center = get_smoke_explosion_center(drone_init_pos, drone_speed, theta, t_drop, t_detonate)
    sink_distance = smoke_sink_speed * (t - detonate_t)
    real_z = max(explosion_center[2] - sink_distance, 10.0)  # 避免落地
    return np.array([explosion_center[0], explosion_center[1], real_z])


# 1.5 遮蔽判定（文档1-8“烟幕阻挡导弹视线”）
def is_segment_sphere_intersect(missile_pos, vertex_pos, smoke_center, smoke_radius):
    vec_AC = smoke_center - missile_pos
    vec_AB = vertex_pos - missile_pos
    ab_sq_mag = np.dot(vec_AB, vec_AB)
    if ab_sq_mag < 1e-12:  # 避免线段长度为0的计算误差
        return False
    k = np.dot(vec_AC, vec_AB) / ab_sq_mag
    # 计算线段到球心的最短距离
    if k <= 0.0:
        min_dist = np.linalg.norm(vec_AC)
    elif 0.0 < k < 1.0:
        cross_prod = np.cross(vec_AC, vec_AB)
        min_dist = np.linalg.norm(cross_prod) / np.sqrt(ab_sq_mag)
    else:
        min_dist = np.linalg.norm(smoke_center - vertex_pos)
    return min_dist <= smoke_radius + 1e-6  # 允许微小计算误差


# 1.6 单枚弹遮蔽判定（文档1-5“起爆后20s有效”）
def is_effective_occlusion(t, drone_speed, theta, t_drop, t_detonate):
    detonate_t = t_drop + t_detonate
    smoke_start = detonate_t
    smoke_end = detonate_t + smoke_valid_time
    # 时间有效性判定
    if t < smoke_start - 1e-6 or t > smoke_end + 1e-6 or t > missile_time_to_fake:
        return False
    # 位置有效性判定
    missile_pos = missile_real_time_position(missile_M1, t)
    smoke_center = smoke_real_time_center(drone_FY1, drone_speed, theta, t_drop, t_detonate, t)
    if np.all(smoke_center == 0):  # 起爆前烟幕无效
        return False
    # 任一顶点视线被挡即有效（文档1-5“云团中心10m范围有效”）
    for vertex in target_vertices:
        if is_segment_sphere_intersect(missile_pos, vertex, smoke_center, smoke_radius):
            return True
    return False


# -------------------------- 2. 优化1：约束处理（柔性间隔+不落地约束，文档1-5） --------------------------
def adjust_t_drop(t_drop_list):
    """柔性调整投放间隔，满足≥1s，保留原参数趋势（优化核心）"""
    adjusted = sorted(t_drop_list)
    for i in range(1, len(adjusted)):
        if adjusted[i] - adjusted[i - 1] < 1.0:
            adjusted[i] = adjusted[i - 1] + 1.0  # 最小化扰动，避免破坏参数关联
    return adjusted


# 2.1 交叉后约束修复（柔性调整投放间隔）
def on_crossover(ga_instance, offspring_crossover):
    for candidate in offspring_crossover:
        # 提取3枚弹的投放时刻（每枚弹第3个基因：索引4i+2）
        t_drop_list = [candidate[i * 4 + 2] for i in range(NUM_SMOKE)]
        adjusted_t_drop = adjust_t_drop(t_drop_list)
        # 回写调整后的投放时刻（保留弹间参数关联）
        for i in range(NUM_SMOKE):
            candidate[i * 4 + 2] = adjusted_t_drop[i]
    return offspring_crossover


# 2.2 变异后约束修复（同交叉修复逻辑）
def on_mutation(ga_instance, offspring_mutation):
    for candidate in offspring_mutation:
        t_drop_list = [candidate[i * 4 + 2] for i in range(NUM_SMOKE)]
        adjusted_t_drop = adjust_t_drop(t_drop_list)
        for i in range(NUM_SMOKE):
            candidate[i * 4 + 2] = adjusted_t_drop[i]
    return offspring_mutation


# 2.3 生成有效初始种群（预过滤间隔约束+位置约束，提升效率）
def generate_valid_population(pop_size, num_genes, gene_space, num_smoke):
    population = []
    while len(population) < pop_size:
        candidate = []
        valid_candidate = True
        for i in range(num_smoke):
            # 按基因范围随机生成单枚弹参数
            speed = np.random.uniform(gene_space[i * 4]['low'], gene_space[i * 4]['high'])
            theta = np.random.uniform(gene_space[i * 4 + 1]['low'], gene_space[i * 4 + 1]['high'])
            t_drop = np.random.uniform(gene_space[i * 4 + 2]['low'], gene_space[i * 4 + 2]['high'])
            t_detonate = np.random.uniform(gene_space[i * 4 + 3]['low'], gene_space[i * 4 + 3]['high'])

            # 预筛选1：烟幕20s内不落地（文档1-5隐含约束）
            explosion_z = 1800 - 0.5 * g * (t_detonate ** 2)
            if explosion_z - 3 * 20 < 10:  # 20s后z≥10
                valid_candidate = False
                break

            # 预筛选2：烟幕在导弹-真目标之间（y∈[-200,400]，z∈[800,1200]，文档1-8）
            t_det = t_drop + t_detonate
            explosion_y = drone_FY1[1] + speed * math.sin(theta) * t_det
            if explosion_y < -200 or explosion_y > 400 or explosion_z < 800 or explosion_z > 1200:
                valid_candidate = False
                break

            candidate.extend([speed, theta, t_drop, t_detonate])

        # 预筛选3：投放间隔≥1s（文档1-5）
        if valid_candidate:
            t_drop_list = [candidate[i * 4 + 2] for i in range(num_smoke)]
            adjusted_t_drop = adjust_t_drop(t_drop_list)
            # 检查调整后是否仍在基因范围（避免超出投放时刻约束）
            for i in range(num_smoke):
                if adjusted_t_drop[i] > gene_space[i * 4 + 2]['high']:
                    valid_candidate = False
                    break
            if valid_candidate:
                # 回写调整后的投放时刻
                for i in range(num_smoke):
                    candidate[i * 4 + 2] = adjusted_t_drop[i]
                population.append(candidate)
    return np.array(population)


# -------------------------- 3. 优化2：适应度函数（动态步长+位置预筛，文档1-8） --------------------------
def calculate_single_duration(drone_speed, theta, t_drop, t_detonate, time_step=0.01):
    """计算单枚弹的有效遮蔽时长（用于结果输出）"""
    detonate_t = t_drop + t_detonate
    smoke_start = detonate_t
    smoke_end = detonate_t + smoke_valid_time
    total_duration = 0.0
    prev_effective = False
    t = smoke_start
    while t <= smoke_end + 1e-6 and t <= missile_time_to_fake:
        current_effective = is_effective_occlusion(t, drone_speed, theta, t_drop, t_detonate)
        if current_effective and not prev_effective:
            start_t = t
        elif not current_effective and prev_effective:
            total_duration += t - start_t
        prev_effective = current_effective
        t += time_step
    if prev_effective:
        total_duration += min(smoke_end, missile_time_to_fake) - start_t
    return round(total_duration, 3)


def calculate_joint_duration(params):
    """优化核心：动态步长计算联合遮蔽时长（平衡精度与效率）"""
    # 1. 获取每枚弹的有效时间区间
    smoke_intervals = []
    for param in params:
        t_det = param[2] + param[3]
        smoke_intervals.append((t_det, t_det + smoke_valid_time))
    # 2. 联合时间范围（覆盖所有弹的有效区间）
    smoke_start = min([s[0] for s in smoke_intervals])
    smoke_end = max([s[1] for s in smoke_intervals])
    smoke_end = min(smoke_end, missile_time_to_fake)  # 不超过导弹命中时间
    if smoke_start >= smoke_end:
        return 0.0

    total_duration = 0.0
    prev_effective = False
    t = smoke_start

    while t <= smoke_end + 1e-6:
        # 动态步长：有效区间内高精度（0.01s），否则低精度（0.1s）
        in_any_smoke = any([s[0] - 1e-6 <= t <= s[1] + 1e-6 for s in smoke_intervals])
        step = 0.01 if in_any_smoke else 0.1

        # 联合遮蔽判定：任一弹有效即有效（文档1-8“遮蔽可不连续”）
        current_effective = False
        for param in params:
            if is_effective_occlusion(t, *param):
                current_effective = True
                break

        # 累加有效时长
        if current_effective and not prev_effective:
            start_t = t
        elif not current_effective and prev_effective:
            total_duration += t - start_t

        prev_effective = current_effective
        t += step

    # 处理区间末尾仍有效的情况
    if prev_effective:
        total_duration += smoke_end - start_t
    return round(total_duration, 3)


def fitness_func(ga_instance, solution, solution_idx):
    """适应度函数：联合遮蔽时长（优化核心目标，文档1-8）"""
    params = []
    for i in range(NUM_SMOKE):
        base = i * 4
        drone_speed = solution[base]
        theta = solution[base + 1]
        t_drop = solution[base + 2]
        t_detonate = solution[base + 3]

        # 双重约束检查（避免交叉/变异后出现无效解）
        # 1. 速度约束（文档1-7）
        if not (drone_speed_range[0] <= drone_speed <= drone_speed_range[1]):
            return 0.0
        # 2. 烟幕不落地约束（文档1-5隐含）
        explosion_z = 1800 - 0.5 * g * (t_detonate ** 2)
        if explosion_z - 3 * 20 < 10:
            return 0.0
        # 3. 烟幕位置约束（文档1-8“导弹-真目标之间”）
        t_det = t_drop + t_detonate
        explosion_y = drone_FY1[1] + drone_speed * math.sin(theta) * t_det
        if explosion_y < -200 or explosion_y > 400 or explosion_z < 800 or explosion_z > 1200:
            return 0.0

        params.append((drone_speed, theta, t_drop, t_detonate))

    # 计算联合遮蔽时长（适应度值）
    joint_duration = calculate_joint_duration(params)
    return joint_duration


# -------------------------- 4. 优化3：GA参数适配（12维变量，文档1-11） --------------------------
def custom_crossover(ga_instance, parents, offspring_size, ga_rng):
    """自定义分段交叉：仅在弹间交叉，保留单弹参数关联性（优化核心）"""
    offspring = []
    num_parents = len(parents)
    for i in range(offspring_size[0]):
        # 随机选择两个父代
        parent1_idx = ga_rng.choice(num_parents)
        parent2_idx = ga_rng.choice(num_parents)
        while parent1_idx == parent2_idx:  # 避免父代相同
            parent2_idx = ga_rng.choice(num_parents)
        parent1 = parents[parent1_idx]
        parent2 = parents[parent2_idx]

        # 分段交叉点：仅在弹与弹之间（第4、8基因后，对应3枚弹的分割）
        cross_points = [4, 8]
        cross_point = ga_rng.choice(cross_points)

        # 生成子代：前cross_point个基因来自父1，后来自父2
        child = np.concatenate([parent1[:cross_point], parent2[cross_point:]])
        offspring.append(child)
    return np.array(offspring)


# 4.1 每代进度回调
def on_generation(ga_instance):
    best_fitness = np.max(ga_instance.last_generation_fitness)
    print(f"当前迭代次数：{ga_instance.generations_completed}, 最大联合遮蔽时长：{best_fitness:.2f}s")


# -------------------------- 5. 遗传算法主程序（整合所有优化） --------------------------
if __name__ == "__main__":
    # 定义基因范围（3枚弹×4参数，共12个基因，文档1-7、1-8）
    gene_space = []
    for i in range(NUM_SMOKE):
        gene_space += [
            {'low': 70, 'high': 140},  # 基因1：无人机速度（文档1-7）
            {'low': 0, 'high': 2 * math.pi},  # 基因2：航向角（0-360°）
            {'low': 0, 'high': 30},  # 基因3：投放延迟（≤导弹命中前）
            {'low': 0, 'high': 10}  # 基因4：起爆延迟（确保不落地）
        ]

    # 生成有效初始种群（优化3：种群规模80，适配12维变量）
    initial_population = generate_valid_population(
        pop_size=80,
        num_genes=4 * NUM_SMOKE,
        gene_space=gene_space,
        num_smoke=NUM_SMOKE
    )
    print(f"初始有效种群生成完成（规模：{len(initial_population)}）")

    # 初始化GA（适配K_tournament参数）
    ga_instance = pygad.GA(
        num_generations=120,  # 迭代次数120（平衡效率与搜索充分性）
        num_parents_mating=20,  # 父代数量20（提升子代多样性）
        fitness_func=fitness_func,  # 适应度函数（联合遮蔽时长）
        sol_per_pop=80,  # 种群规模80（适配12维变量）
        num_genes=4 * NUM_SMOKE,  # 基因数量12（3枚弹×4参数）
        gene_space=gene_space,  # 基因范围约束
        parent_selection_type="tournament",  # 选择策略：锦标赛（避免局部最优）
        K_tournament=3,  # 锦标赛规模（匹配提供的__init__参数）
        crossover_type=custom_crossover,  # 交叉策略：自定义分段交叉（优化核心）
        mutation_type="random",  # 变异策略：随机变异
        mutation_percent_genes=5,  # 变异基因比例5%（控制扰动幅度）
        mutation_num_genes=1,  # 每子代变异1个基因（避免无效解）
        parallel_processing=8,  # 并行计算（提升效率）
        on_generation=on_generation,  # 每代回调（监控进度）
        on_crossover=on_crossover,  # 交叉后约束修复
        on_mutation=on_mutation,  # 变异后约束修复
        initial_population=initial_population  # 自定义有效初始种群
    )

    # 运行GA
    print("开始3枚烟幕弹联合优化（基于《A题.pdf》约束）...")
    ga_instance.run()

    # 解析最优解（文档1-11“输出投放策略”）
    solution, solution_fitness, solution_idx = ga_instance.best_solution()
    single_durations = []  # 单枚弹遮蔽时长（用于结果输出）
    for i in range(NUM_SMOKE):
        base = i * 4
        param_i = solution[base:base + 4]
        single_durations.append(calculate_single_duration(*param_i))

    # 输出最优策略（符合文档1-8“投放策略包含飞行/投放/起爆参数”）
    print("\n" + "=" * 60)
    print("《A题.pdf》第三题最优投放策略（FY1投放3枚烟幕弹）")
    print("=" * 60)
    for i in range(NUM_SMOKE):
        base = i * 4
        drone_speed = solution[base]
        theta = solution[base + 1]
        t_drop = solution[base + 2]
        t_detonate = solution[base + 3]
        t_det = t_drop + t_detonate

        # 计算投放点与起爆点（文档1-8要求）
        drop_x = drone_FY1[0] + drone_speed * math.cos(theta) * t_drop
        drop_y = drone_FY1[1] + drone_speed * math.sin(theta) * t_drop
        drop_z = drone_FY1[2]
        explosion_point = get_smoke_explosion_center(
            drone_FY1, drone_speed, theta, t_drop, t_detonate
        )
        theta_deg = (theta * 180 / math.pi) % 360  # 弧度转角度（便于阅读）

        # 输出单弹策略
        print(f"\n【烟幕弹{i + 1}】")
        print(f"1. 无人机参数：速度={drone_speed:.2f}m/s，航向角={theta_deg:.2f}°")
        print(f"2. 时间参数：投放时刻={t_drop:.2f}s，起爆间隔={t_detonate:.2f}s，起爆时刻={t_det:.2f}s")
        print(f"3. 空间参数：投放点=({drop_x:.2f}, {drop_y:.2f}, {drop_z:.2f})m")
        print(f"             起爆点=({explosion_point[0]:.2f}, {explosion_point[1]:.2f}, {explosion_point[2]:.2f})m")
        print(f"4. 单弹有效遮蔽时长：{single_durations[i]:.2f}s")

    # 输出联合结果（文档1-8“最大化联合遮蔽时长”）
    print(f"\n【联合优化结果】")
    print(f"最大联合有效遮蔽时长：{solution_fitness:.2f}s")
    print("=" * 60)
